import os
from flask import Flask, request, jsonify, send_from_directory
import torch
from transformers import BertTokenizer
from nete import Model
 
# === 自动定位项目根目录（web_test/） ===
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
# === 设置静态文件路径 ===
app = Flask(__name__, static_folder=os.path.join(BASE_DIR, "static"))
 
# === 设备设置 ===
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"✅ 当前设备：{DEVICE}")
 
# === 路径设置 ===
model_path = os.path.join(BASE_DIR, "params", "best_bert.pth")
tokenizer_path = os.path.join(
    BASE_DIR,
    "model",
    "bert-base-chinese",
    "models--bert-base-chinese",
    "snapshots",
    "c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
)
 
# === 模型加载 ===
print("🚀 [1] 正在加载模型...")
model = Model(dropout_rate=0.3, unfreeze_last_n_layers=2).to(DEVICE)
model.load_state_dict(torch.load(model_path, map_location=DEVICE))
model.eval()
print("✅ [1] 模型加载完成")
 
# === 分词器加载 ===
print("🚀 [2] 正在加载 tokenizer...")
tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
print("✅ [2] tokenizer 加载完成")
 
# === 类别标签 ===
labels = ["非常差", "差", "一般", "中性", "还行", "满意", "很好", "非常好"]
 
# === 预测接口 ===
@app.route("/predict", methods=["POST"])
def predict():
    data = request.get_json()
    text = data.get("text", "")
    if not text:
        return jsonify({"error": "请输入内容"}), 400
 
    encoded = tokenizer.batch_encode_plus(
        [text],
        truncation=True,
        max_length=128,
        padding="max_length",
        return_tensors="pt",
        return_length=True
    )
 
    input_ids = encoded["input_ids"].to(DEVICE)
    attention_mask = encoded["attention_mask"].to(DEVICE)
    token_type_ids = encoded["token_type_ids"].to(DEVICE)
 
    with torch.no_grad():
        output = model(input_ids, attention_mask, token_type_ids)
        pred = torch.argmax(output, dim=1).item()
 
    return jsonify({"label": labels[pred]})
 
# === 首页路由 ===
@app.route("/")
def index():
    return send_from_directory(app.static_folder, "index.html")
 
# === 启动服务 ===
if __name__ == '__main__':
    print("🚀 [3] 启动 Flask 服务中...")
    app.run(debug=True, port=5000)